package vip.liux.front.infrastructure.config;

import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.boot.task.ThreadPoolTaskExecutorCustomizer;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.EnableAsync;
import org.springframework.scheduling.annotation.EnableScheduling;
import org.springframework.scheduling.annotation.SchedulingConfigurer;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.config.ScheduledTaskRegistrar;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;

import static vip.liux.core.thread.CustomizeThreadPoolExecutor.DEFAULT_REJECTED_HANDLER;


@Configuration(proxyBeanMethods = false)
@EnableAsync
@EnableScheduling
public class TaskConfiguration implements ThreadPoolTaskExecutorCustomizer, SchedulingConfigurer {

    private static final Logger log = LoggerFactory.getLogger(TaskConfiguration.class);

    @Override
    public void customize(ThreadPoolTaskExecutor taskExecutor) {
        // 线程池对拒绝任务(无线程可用)的处理策略
        taskExecutor.setRejectedExecutionHandler(DEFAULT_REJECTED_HANDLER);
    }

    @Bean
    public TaskDecorator taskDecorator() {
        return runnable -> {
            try {
                // Right now: Web thread context !
                // Grab the current thread MDC data
                Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
                //获取父线程的context
                RequestAttributes context = RequestContextHolder.currentRequestAttributes();
                return () -> {
                    // Right now: @Async thread context !
                    // Restore the Web thread context's MDC data
                    MDC.setContextMap(copyOfContextMap);
                    //将父线程的context设置进子线程里
                    RequestContextHolder.setRequestAttributes(context);
                    runnable.run();
                };
            } catch (Exception e) {
                return runnable;
            }
        };
    }

    @Override
    public void configureTasks(@NotNull ScheduledTaskRegistrar taskRegistrar) {
        log.debug("Configuring scheduled tasks: {}", taskRegistrar);
    }
}
